package sample;

import org.objectweb.asm.*;
import org.objectweb.asm.tree.*;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.tree.analysis.*;

import java.util.ArrayList;
import java.util.List;

import static org.objectweb.asm.Opcodes.ILOAD;

public class InstrumentMethod {

    static final String STACK_NAME = "sample.Stack".replace('.', '/');
    static final String FIBER_NAME = "sample.Fiber".replace('.', '/');

    MethodNode methodNode;
    final Frame[] frames;
    final int lvarStack; // ref to Stack
    final int lvarState; // ref to Stack
    final int firstLocal;


    class BlockCode {
        Label startLabel;
        Label endLabel;
        InsnList instructions;
        int start;
        int end;
        int index;

        BlockCode(InsnList instructions, int start, int index) {
            this.instructions = instructions;

            this.start = start;
            this.index = index;
            startLabel = new Label();
            endLabel = new Label();
        }


        void accept(final MethodVisitor mv) {
            mv.visitLabel(startLabel);
            emitRestoreState(mv, frames[start]);

            for (int i = start; i < end; i++) {
                final AbstractInsnNode in = instructions.get(i);
                in.accept(mv);
            }
            if (index < 5) {
                mv.visitInsn(Opcodes.ICONST_0 + index + 1);

            } else {
                mv.visitIntInsn(Opcodes.BIPUSH, index + 1);

            }

            mv.visitMethodInsn(Opcodes.INVOKESTATIC, FIBER_NAME, "saveEntry", "(I)V", false);
            if (end < instructions.size()) {
                emitStoreState(mv, frames[end]);

            }
            mv.visitLabel(endLabel);
            mv.visitInsn(Opcodes.RETURN);
        }

    }

    public InstrumentMethod(MethodNode methodNode) {
        this.methodNode = methodNode;
        this.lvarState = methodNode.maxLocals++;
        this.lvarStack = methodNode.maxLocals++;
        methodNode.maxStack++;
        methodNode.maxStack++;

        this.firstLocal = ((methodNode.access & Opcodes.ACC_STATIC) == Opcodes.ACC_STATIC) ? 0 : 1;

        Analyzer<BasicValue> analyzer = new Analyzer<>(new BasicInterpreter());
        try {
            analyzer.analyze(methodNode.name, methodNode);
            frames = analyzer.getFrames();
        } catch (AnalyzerException e) {
            throw new RuntimeException(e);
        }

    }

    static void accept(final AnnotationVisitor av, final String name,
                       final Object value) {
        if (av != null) {
            if (value instanceof String[]) {
                String[] typeconst = (String[]) value;
                av.visitEnum(name, typeconst[0], typeconst[1]);
            } else if (value instanceof AnnotationNode) {
                AnnotationNode an = (AnnotationNode) value;
                an.accept(av.visitAnnotation(name, an.desc));
            } else if (value instanceof List) {
                AnnotationVisitor v = av.visitArray(name);
                if (v != null) {
                    List<?> array = (List<?>) value;
                    for (int j = 0; j < array.size(); ++j) {
                        accept(v, null, array.get(j));
                    }
                    v.visitEnd();
                }
            } else {
                av.visit(name, value);
            }
        }
    }

    void accept(final MethodVisitor mv) {
        // visits the method parameters
        List<ParameterNode> parameters = methodNode.parameters;
        Object annotationDefault = methodNode.annotationDefault;
        List<AnnotationNode> visibleAnnotations = methodNode.visibleAnnotations;
        List<AnnotationNode> invisibleAnnotations = methodNode.invisibleAnnotations;
        List<TypeAnnotationNode> visibleTypeAnnotations = methodNode.visibleTypeAnnotations;
        List<TypeAnnotationNode> invisibleTypeAnnotations = methodNode.invisibleTypeAnnotations;
        List<AnnotationNode>[] visibleParameterAnnotations = methodNode.visibleParameterAnnotations;
        List<AnnotationNode>[] invisibleParameterAnnotations = methodNode.invisibleParameterAnnotations;
        List<Attribute> attrs = methodNode.attrs;
        InsnList instructions = methodNode.instructions;
        List<TryCatchBlockNode> tryCatchBlocks = methodNode.tryCatchBlocks;
        List<LocalVariableNode> localVariables = methodNode.localVariables;
        List<LocalVariableAnnotationNode> visibleLocalVariableAnnotations = methodNode.visibleLocalVariableAnnotations;
        List<LocalVariableAnnotationNode> invisibleLocalVariableAnnotations = methodNode.invisibleLocalVariableAnnotations;
        int maxStack = methodNode.maxStack;

        int maxLocals = methodNode.maxLocals;

        int i, j, n;
        n = parameters == null ? 0 : parameters.size();
        for (i = 0; i < n; i++) {
            ParameterNode parameter = parameters.get(i);
            mv.visitParameter(parameter.name, parameter.access);
        }
        // visits the method attributes
        if (annotationDefault != null) {
            AnnotationVisitor av = mv.visitAnnotationDefault();
            accept(av, null, annotationDefault);
            if (av != null) {
                av.visitEnd();
            }
        }
        n = visibleAnnotations == null ? 0 : visibleAnnotations.size();
        for (i = 0; i < n; ++i) {
            AnnotationNode an = visibleAnnotations.get(i);
            an.accept(mv.visitAnnotation(an.desc, true));
        }
        n = invisibleAnnotations == null ? 0 : invisibleAnnotations.size();
        for (i = 0; i < n; ++i) {
            AnnotationNode an = invisibleAnnotations.get(i);
            an.accept(mv.visitAnnotation(an.desc, false));
        }
        n = visibleTypeAnnotations == null ? 0 : visibleTypeAnnotations.size();
        for (i = 0; i < n; ++i) {
            TypeAnnotationNode an = visibleTypeAnnotations.get(i);
            an.accept(mv.visitTypeAnnotation(an.typeRef, an.typePath, an.desc, true));
        }
        n = invisibleTypeAnnotations == null ? 0 : invisibleTypeAnnotations.size();
        for (i = 0; i < n; ++i) {
            TypeAnnotationNode an = invisibleTypeAnnotations.get(i);
            an.accept(mv.visitTypeAnnotation(an.typeRef, an.typePath, an.desc, false));
        }
        n = visibleParameterAnnotations == null ? 0 : visibleParameterAnnotations.length;
        for (i = 0; i < n; ++i) {
            List<?> l = visibleParameterAnnotations[i];
            if (l == null) {
                continue;
            }
            for (j = 0; j < l.size(); ++j) {
                AnnotationNode an = (AnnotationNode) l.get(j);
                an.accept(mv.visitParameterAnnotation(i, an.desc, true));
            }
        }
        n = invisibleParameterAnnotations == null ? 0 : invisibleParameterAnnotations.length;
        for (i = 0; i < n; ++i) {
            List<?> l = invisibleParameterAnnotations[i];
            if (l == null) {
                continue;
            }
            for (j = 0; j < l.size(); ++j) {
                AnnotationNode an = (AnnotationNode) l.get(j);
                an.accept(mv.visitParameterAnnotation(i, an.desc, false));
            }
        }
        n = attrs == null ? 0 : attrs.size();
        for (i = 0; i < n; ++i) {
            mv.visitAttribute(attrs.get(i));
        }
        // visits the method's code
        if (instructions.size() > 0) {
            mv.visitCode();
            // visits try catch blocks
            n = tryCatchBlocks == null ? 0 : tryCatchBlocks.size();
            for (i = 0; i < n; ++i) {
                tryCatchBlocks.get(i).updateIndex(i);
                tryCatchBlocks.get(i).accept(mv);
            }

            // visits instructions
            final int numIns = instructions.size();
            Label start = null;
            Label end = null;
            List<BlockCode> blockCodes = new ArrayList<>();
            BlockCode blockCode = new BlockCode(instructions, 0, 0);
            blockCodes.add(blockCode);
            for (i = 0; i < numIns; i++) {
                boolean yield = false;
                final AbstractInsnNode in = instructions.get(i);
                if (in instanceof MethodInsnNode) {
                    MethodInsnNode methodInsnNode = (MethodInsnNode) in;
                    yield = (methodInsnNode.owner.equals(FIBER_NAME) && methodInsnNode.name.equals("yield"));

                }
                if (yield) {
                    blockCode.end = i;
                    blockCode = new BlockCode(instructions, i + 1, blockCodes.size());
                    blockCodes.add(blockCode);


                }
                if (in instanceof LabelNode) {
                    LabelNode labelNode = (LabelNode) in;
                    if (start == null) {
                        start = labelNode.getLabel();
                    }
                    end = labelNode.getLabel();

                }
            }
            blockCode.end = numIns;

            Label[] lMethodCalls = new Label[blockCodes.size()];
            for (i = 0; i < blockCodes.size(); i++) {
                blockCode = blockCodes.get(i);
                lMethodCalls[i] = blockCode.startLabel;
            }


            mv.visitFieldInsn(Opcodes.GETSTATIC, FIBER_NAME, "entry", "I");
            mv.visitVarInsn(Opcodes.ISTORE, lvarState);
            mv.visitFieldInsn(Opcodes.GETSTATIC, FIBER_NAME, "stack", "L" + STACK_NAME + ";");
            mv.visitVarInsn(Opcodes.ASTORE, lvarStack);
            mv.visitVarInsn(ILOAD, lvarState);
            mv.visitTableSwitchInsn(0, blockCodes.size() - 1, blockCode.endLabel, lMethodCalls);
            for (BlockCode b : blockCodes) {
                b.accept(mv);
            }

            // visits local variables
            n = localVariables == null ? 0 : localVariables.size();
            for (i = 0; i < n; ++i) {
                localVariables.get(i).accept(mv);
            }

            // visits local variable annotations
            n = visibleLocalVariableAnnotations == null ? 0 : visibleLocalVariableAnnotations.size();
            for (i = 0; i < n; ++i) {
                visibleLocalVariableAnnotations.get(i).accept(mv, true);
            }
            n = invisibleLocalVariableAnnotations == null ? 0 : invisibleLocalVariableAnnotations.size();
            for (i = 0; i < n; ++i) {
                invisibleLocalVariableAnnotations.get(i).accept(mv, false);
            }
            // visits maxs
            mv.visitMaxs(maxStack, maxLocals);
        }
        mv.visitEnd();
    }

    private void emitStoreState(MethodVisitor mv, Frame f) {

        if (f == null) {
            return;
        }
        // store operand stack
        for (int i = f.getStackSize(); i-- > 0; ) {
            BasicValue v = (BasicValue) f.getStack(i);
            if (!isNullType(v)) {
                emitStoreValue(mv, v, lvarStack);
            } else {
                mv.visitInsn(Opcodes.POP);
            }

        }

        // store local vars
        for (int i = firstLocal; i < f.getLocals(); i++) {
            BasicValue v = (BasicValue) f.getLocal(i);
            if (!isNullType(v)) {
                mv.visitVarInsn(v.getType().getOpcode(Opcodes.ILOAD), i);
                emitStoreValue(mv, v, lvarStack);
            }
        }


    }


    static void emitStoreValue(MethodVisitor mv, BasicValue v, int lvarStack) throws InternalError, IndexOutOfBoundsException {
        String desc;

        switch (v.getType().getSort()) {
            case Type.OBJECT:
            case Type.ARRAY:
                desc = "(Ljava/lang/Object;L" + STACK_NAME + ";)V";
                break;
            case Type.BOOLEAN:
            case Type.BYTE:
            case Type.SHORT:
            case Type.CHAR:
            case Type.INT:
                desc = "(IL" + STACK_NAME + ";)V";
                break;
            case Type.FLOAT:
                desc = "(FL" + STACK_NAME + ";)V";
                break;
            case Type.LONG:
                desc = "(JL" + STACK_NAME + ";)V";
                break;
            case Type.DOUBLE:
                desc = "(DL" + STACK_NAME + ";)V";
                break;
            default:
                throw new InternalError("Unexpected type: " + v.getType());
        }

        mv.visitVarInsn(Opcodes.ALOAD, lvarStack);
        mv.visitMethodInsn(Opcodes.INVOKESTATIC, STACK_NAME, "push", desc, false);
    }

    static void emitConst(MethodVisitor mv, int value) {
        if (value >= -1 && value <= 5)
            mv.visitInsn(Opcodes.ICONST_0 + value);
        else if ((byte) value == value)
            mv.visitIntInsn(Opcodes.BIPUSH, value);
        else if ((short) value == value)
            mv.visitIntInsn(Opcodes.SIPUSH, value);
        else
            mv.visitLdcInsn(value);
    }

    static boolean isNullType(BasicValue v) {
        return (v == BasicValue.UNINITIALIZED_VALUE) || (v.isReference() && v.getType().getInternalName().equals("null"));
    }

    void emitRestoreState(MethodVisitor mv, Frame f) {
        if (f == null) {
            return;
        }

        // restore local vars
        for (int i = firstLocal; i < f.getLocals(); i++) {
            BasicValue v = (BasicValue) f.getLocal(i);
            if (!isNullType(v)) {
                emitRestoreValue(mv, v, lvarStack);
                mv.visitVarInsn(v.getType().getOpcode(Opcodes.ISTORE), i);
            } else if (v != BasicValue.UNINITIALIZED_VALUE) {
                mv.visitInsn(Opcodes.ACONST_NULL);
                mv.visitVarInsn(Opcodes.ASTORE, i);
            }
        }

        // restore operand stack
        for (int i = 0; i < f.getStackSize(); i++) {
            BasicValue v = (BasicValue) f.getStack(i);
            if (!isNullType(v)) {
                emitRestoreValue(mv, v, lvarStack);
            } else {
                mv.visitInsn(Opcodes.ACONST_NULL);
            }
        }

    }

    static void emitRestoreValue(MethodVisitor mv, BasicValue v, int lvarStack) {
        mv.visitVarInsn(Opcodes.ALOAD, lvarStack);

        switch (v.getType().getSort()) {
            case Type.OBJECT:
                String internalName = v.getType().getInternalName();
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popObject", "()Ljava/lang/Object;", false);
                if (!internalName.equals("java/lang/Object"))  // don't cast to Object ;)
                    mv.visitTypeInsn(Opcodes.CHECKCAST, internalName);
                break;
            case Type.ARRAY:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popObject", "()Ljava/lang/Object;", false);
                mv.visitTypeInsn(Opcodes.CHECKCAST, v.getType().getDescriptor());
                break;
            case Type.BYTE:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popInt", "()I", false);
                mv.visitInsn(Opcodes.I2B);
                break;
            case Type.SHORT:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popInt", "()I", false);
                mv.visitInsn(Opcodes.I2S);
                break;
            case Type.CHAR:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popInt", "()I", false);
                mv.visitInsn(Opcodes.I2C);
                break;
            case Type.BOOLEAN:
            case Type.INT:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popInt", "()I", false);
                break;
            case Type.FLOAT:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popFloat", "()F", false);
                break;
            case Type.LONG:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popLong", "()J", false);
                break;
            case Type.DOUBLE:
                mv.visitMethodInsn(Opcodes.INVOKEVIRTUAL, STACK_NAME, "popDouble", "()D", false);
                break;
            default:
                throw new InternalError("Unexpected type: " + v.getType());
        }
    }

}
